package com.maben.aspect;

import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.MDC;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.UUID;

/**
 * AOP添加MDC的内容
 */
@Aspect
@Component
public class TraceIdAspect {
    public static final String TRACE_ID_HEADER = "X-Trace-Id";
    public static final String TRACE_ID_MDC_KEY = "traceId";
    /**
     * 定义切入点：接口
     */
    @Pointcut("@annotation(org.springframework.web.bind.annotation.GetMapping) || " +
            "@annotation(org.springframework.web.bind.annotation.RequestMapping) || " +
            "@annotation(org.springframework.web.bind.annotation.PostMapping) || " +
            "@annotation(org.springframework.web.bind.annotation.PutMapping) || " +
            "@annotation(org.springframework.web.bind.annotation.DeleteMapping)")
    public void restEndpoint() {}

    /**
     * 切入点：非接口
     */
    @Pointcut("@annotation(com.maben.aspect.Traceable)")
    public void traceableMethod() {}

    /**
     * 前置方法
     */
    @Before("restEndpoint() || traceableMethod()")
    public void before(JoinPoint joinPoint) {
        String traceId = null;

        // 尝试从HTTP请求头获取traceId
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes != null) {
            HttpServletRequest request = attributes.getRequest();
            traceId = request.getHeader(TRACE_ID_HEADER);
        }

        // 如果没有从请求头获取到traceId，则生成一个新的
        if (traceId == null || traceId.isEmpty()) {
            traceId = UUID.randomUUID().toString().replace("-", "");
        }

        // 将traceId放入MDC
        MDC.put(TRACE_ID_MDC_KEY, traceId);
    }

    /**
     * 后置方法
     */
    @After("restEndpoint() || traceableMethod()")
    public void after(JoinPoint joinPoint) {
        // 清除MDC中的traceId，防止内存泄漏
        MDC.remove(TRACE_ID_MDC_KEY);
    }

}